import torch.nn as nn
from collections import OrderedDict
import torch

class eca_layer(nn.Module):
    """Constructs a ECA module.
    Args:
        channel: Number of channels of the input feature map
        k_size: Adaptive selection of kernel size
    """
    def __init__(self, channel, k_size=3):
        super(eca_layer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        # x: input features with shape [b, c, h, w]
        b, c, h, w = x.size()

        # feature descriptor on the global spatial information
        y = self.avg_pool(x)

        # Two different branches of ECA module
        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)

        # Multi-scale information fusion
        y = self.sigmoid(y)

        return x * y.expand_as(x)

def conv_layer(in_channels, out_channels, kernel_size, stride=1, dilation=1, groups=1):
    padding = int((kernel_size - 1) / 2) * dilation
    return nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, bias=True, dilation=dilation,
                     groups=groups)


def norm(norm_type, nc):
    norm_type = norm_type.lower()
    if norm_type == 'batch':
        layer = nn.BatchNorm2d(nc, affine=True)
    elif norm_type == 'instance':
        layer = nn.InstanceNorm2d(nc, affine=False)
    else:
        raise NotImplementedError('normalization layer [{:s}] is not found'.format(norm_type))
    return layer


def pad(pad_type, padding):
    pad_type = pad_type.lower()
    if padding == 0:
        return None
    if pad_type == 'reflect':
        layer = nn.ReflectionPad2d(padding)
    elif pad_type == 'replicate':
        layer = nn.ReplicationPad2d(padding)
    else:
        raise NotImplementedError('padding layer [{:s}] is not implemented'.format(pad_type))
    return layer


def get_valid_padding(kernel_size, dilation):
    kernel_size = kernel_size + (kernel_size - 1) * (dilation - 1)
    padding = (kernel_size - 1) // 2
    return padding


def conv_block(in_nc, out_nc, kernel_size, stride=1, dilation=1, groups=1, bias=True,
               pad_type='zero', norm_type=None, act_type='relu'):
    padding = get_valid_padding(kernel_size, dilation)
    p = pad(pad_type, padding) if pad_type and pad_type != 'zero' else None
    padding = padding if pad_type == 'zero' else 0

    c = nn.Conv2d(in_nc, out_nc, kernel_size=kernel_size, stride=stride, padding=padding,
                  dilation=dilation, bias=bias, groups=groups)
    a = activation(act_type) if act_type else None
    n = norm(norm_type, out_nc) if norm_type else None
    return sequential(p, c, n, a)


def activation(act_type, inplace=True, neg_slope=0.05, n_prelu=1):
    act_type = act_type.lower()
    if act_type == 'relu':
        layer = nn.ReLU(inplace)
    elif act_type == 'lrelu':
        layer = nn.LeakyReLU(neg_slope, inplace) # Two options → 1)  nn.LeakyReLU(neg_slope, inplace) 2) nn.LeakyReLU(inplace)
    elif act_type == 'prelu':
        layer = nn.PReLU(num_parameters=n_prelu, init=neg_slope)
    else:
        raise NotImplementedError('activation layer [{:s}] is not found'.format(act_type))
    return layer


class ShortcutBlock(nn.Module):
    def __init__(self, submodule):
        super(ShortcutBlock, self).__init__()
        self.sub = submodule

    def forward(self, x):
        output = x + self.sub(x)
        return output

def mean_channels(F):
    assert(F.dim() == 4)
    spatial_sum = F.sum(3, keepdim=True).sum(2, keepdim=True)
    return spatial_sum / (F.size(2) * F.size(3))

def stdv_channels(F):
    assert(F.dim() == 4)
    F_mean = mean_channels(F)
    F_variance = (F - F_mean).pow(2).sum(3, keepdim=True).sum(2, keepdim=True) / (F.size(2) * F.size(3))
    return F_variance.pow(0.5)

def sequential(*args):
    if len(args) == 1:
        if isinstance(args[0], OrderedDict):
            raise NotImplementedError('sequential does not support OrderedDict input.')
        return args[0]
    modules = []
    for module in args:
        if isinstance(module, nn.Sequential):
            for submodule in module.children():
                modules.append(submodule)
        elif isinstance(module, nn.Module):
            modules.append(module)
    return nn.Sequential(*modules)

# contrast-aware channel attention module
class CCALayer(nn.Module):
    def __init__(self, channel, reduction=16):
        super(CCALayer, self).__init__()

        self.contrast = stdv_channels
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv_du = nn.Sequential(
            nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True),
            nn.Sigmoid()
        )


    def forward(self, x):
        y = self.contrast(x) + self.avg_pool(x)

        y = self.conv_du(y)
        return x * y


class ESALayer(nn.Module):
    def __init__(self, n_feats):
        super(ESALayer, self).__init__()



        #
        # self.conv00 = conv_block(n_feats, n_feats,3)
        # self.conv01 = conv_block(n_feats, n_feats,3)

        self.reduce_channels = conv_block(n_feats, n_feats//4, 1)
        self.reduce_spatial_size = conv_block(n_feats//4, n_feats//4, 3, stride=2)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//4, n_feats, 1)

        self.conv1 = conv_block(n_feats//4, n_feats//4,3)
        self.conv2 = conv_block(n_feats // 4, n_feats//4, 3)
        self.conv3 = conv_block(n_feats // 4, n_feats//4, 3)

        self.sigmoid = nn.Sigmoid()


        self.upsampling  = nn.PixelShuffle(2)

        self.act = nn.LeakyReLU(inplace=True)

    def forward(self, x):

        # x = self.conv01(self.act(self.conv00(x)))

        rc = self.reduce_channels(x)

        rs = self.reduce_spatial_size(rc)

        pool = self.pool(rs)

        conv = self.conv1(self.conv2(self.conv3(pool)))

        up = self.upsampling(conv)



        out = self.sigmoid(up)

        return out*x



def channel_shuffle(x, groups):
    batchsize, num_channels, height, width = x.data.size()

    channels_per_group = num_channels // groups

    # reshape
    x = x.view(batchsize, groups,
               channels_per_group, height, width)

    # transpose
    # - contiguous() required if transpose() is used before view().
    #   See https://github.com/pytorch/pytorch/issues/764
    x = torch.transpose(x, 1, 2).contiguous()

    # flatten
    x = x.view(batchsize, -1, height, width)

    return x


class ESAB_eca(nn.Module):
    def __init__(self, n_feats, distillation_rate=0.25):
        super(ESAB_eca, self).__init__()



        self.reduce_channels = nn.Conv2d(n_feats, n_feats//4,1)
        self.reduce_spatial_size = nn.Conv2d(n_feats//4, n_feats//4, 3, stride=2, padding=1)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//4, n_feats, 1)

        self.conv1 = conv_block(n_feats//4, n_feats//4,3,dilation=1)
        self.conv2 = conv_block(n_feats // 4, n_feats//4, 3,dilation=2)
        #self.conv3 = conv_block(n_feats // 4, n_feats//4, 3)

       # self.sigmoid = nn.Sigmoid()


        self.conv00 = conv_block(n_feats, n_feats,3)
        self.conv01 = conv_block(n_feats, n_feats,3)
        self.act =  nn.LeakyReLU(inplace=False)


        #self.upsampling  = nn.PixelShuffle(2)


        # ###############
        # in_channels = int(n_feats//4)
        #
        # self.distilled_channels = int(in_channels * distillation_rate)
        # self.remaining_channels = int(in_channels - self.distilled_channels)
        # #self.c1 = conv_layer(in_channels, in_channels, 3)
        # self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        # self.act = nn.LeakyReLU(inplace=True)# activation('lrelu', neg_slope=0.05)
       # self.act2 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)
      #  self.act3 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)

        # self.scale1 = Scale(1)
        # self.scale2 = Scale(1)
        # self.scale3 = Scale(1)
        # self.scale4 = Scale(1)
        ###############
        self.eca = eca_layer(n_feats)



    def forward(self, x):

        x  = self.conv00(self.conv01(x))
        rc = self.reduce_channels(x)

        rs = torch.nn.functional.upsample(rc, 2, model='nearest')
        #rs =self.reduce_spatial_size(rc)


        #up = self.upsampling(x)

        #print("before pool", rs.shape)

        #pool = self.pool(rs)

      #  print("After ", pool.shape)
        conv = self.conv2(rs)
        conv = conv + self.conv1(rs)


      #  print(rc.shape)
       # up = torch.nn.functional.upsample(conv, size=(rc.shape[2],rc.shape[3]), mode='bicubic')
        up = self.pool(conv)
       # print(up.shape)
        up =  up + rc
        out = self.eca(self.increase_channels(up)) * x


        #out_c1 = self.act(self.c1(up))

        # distilled_c1, remaining_c1 = torch.split(up, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c2 = self.act(self.c2(remaining_c1))
        # distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c3 = self.act(self.c3(remaining_c2))
        # distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c4 = self.c4(remaining_c3)

        # F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        # F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))
        #
        # out = self.fusion3(torch.cat([F1, F2],1))





  #       out = self.increase_channels(torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1))
  #       out = self.reduce_spatial_size (out) #+ up
  #      #
  #
  # #       rc = self.reduce_spatial_size(rc)
  #
  #       out = self.sigmoid(out)*x


        return out


class ESAB_MODIFIED_2(nn.Module):
    def __init__(self, n_feats, reduction_factor=4, distillation_rate=0.25):
        super(ESAB_MODIFIED_2, self).__init__()



        self.reduce_channels = nn.Conv2d(n_feats, n_feats//reduction_factor,1)
        self.reduce_spatial_size = nn.Conv2d(n_feats//reduction_factor, n_feats//reduction_factor, 3, stride=2, padding=1)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//reduction_factor, n_feats, 1)

        self.conv1 = conv_block(n_feats//reduction_factor, n_feats//reduction_factor,3,dilation=1,act_type='lrelu')
        self.conv2 = conv_block(n_feats // reduction_factor, n_feats//reduction_factor, 3,dilation=2,act_type='lrelu')
        #self.conv3 = conv_block(n_feats // 4, n_feats//4, 3)

        self.sigmoid = nn.Sigmoid()


        self.conv00 = conv_block(n_feats, n_feats,3, act_type=None)
        self.conv01 = conv_block(n_feats, n_feats,3,act_type='lrelu')



        #self.upsampling  = nn.PixelShuffle(2)


        # ###############
        # in_channels = int(n_feats//4)
        #
        # self.distilled_channels = int(in_channels * distillation_rate)
        # self.remaining_channels = int(in_channels - self.distilled_channels)
        # #self.c1 = conv_layer(in_channels, in_channels, 3)
        # self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        #self.act = activation('lrelu', neg_slope=0.01)
       # self.act2 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)
      #  self.act3 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)

        self.scale1 = Scale(1)
        self.scale2 = Scale(1)
        self.scale3 = Scale(1)
        self.scale4 = Scale(1)
        ###############

        self.bottom11 = conv_block(n_feats,n_feats,1,act_type=None)
        self.bottom11_dw = conv_block(n_feats, n_feats,5, groups=n_feats,act_type=None)

        self.bottom22 = conv_block(n_feats, n_feats, 1, act_type=None)
        self.bottom22_dw = conv_block(n_feats, n_feats, 5, groups=n_feats, act_type='lrelu')




    def forward(self, x):


        x  = self.bottom22_dw(self.bottom22(x))
        rc = self.reduce_channels(x)

        rs = self.reduce_spatial_size(rc)






        pool = self.pool(rs)

      #  print("After ", pool.shape)
        conv = self.conv2(pool)
        conv = conv + self.conv1(pool)


      #  print(rc.shape)

       # print(up.shape)

        up =  torch.nn.functional.upsample(conv, size=(rc.shape[2],rc.shape[3]), mode='nearest')
        up =  up + rc
        out = self.increase_channels(up) * self.sigmoid(self.bottom11_dw(self.bottom11(x)))





        #out_c1 = self.act(self.c1(up))

        # distilled_c1, remaining_c1 = torch.split(up, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c2 = self.act(self.c2(remaining_c1))
        # distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c3 = self.act(self.c3(remaining_c2))
        # distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c4 = self.c4(remaining_c3)

        # F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        # F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))
        #
        # out = self.fusion3(torch.cat([F1, F2],1))





  #       out = self.increase_channels(torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1))
  #       out = self.reduce_spatial_size (out) #+ up
  #      #
  #
  # #       rc = self.reduce_spatial_size(rc)
  #
  #       out = self.sigmoid(out)*x


        return out

class ESAB_MODIFIED_CAT(nn.Module):
    def __init__(self, n_feats, reduction_factor=4, distillation_rate=0.25):
        super(ESAB_MODIFIED_CAT, self).__init__()



        self.reduce_channels = nn.Conv2d(n_feats, n_feats//reduction_factor,1)
        self.reduce_spatial_size = nn.Conv2d(n_feats//reduction_factor, n_feats//reduction_factor, 3, stride=2, padding=1)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//reduction_factor, n_feats, 1)

        self.conv1 = conv_block(n_feats//reduction_factor, n_feats//reduction_factor,3,dilation=1,act_type='lrelu')
        self.conv2 = conv_block(n_feats // reduction_factor, n_feats//reduction_factor, 3,dilation=1,act_type='lrelu')
        self.conv3 = conv_block(n_feats // reduction_factor, n_feats // reduction_factor, 3, dilation=1,act_type='lrelu')
        #self.conv3 = conv_block(n_feats // 4, n_feats//4, 3)

        self.sigmoid = nn.Sigmoid()


        self.conv00 = conv_block(n_feats, n_feats,3, act_type=None)
        self.conv01 = conv_block(n_feats, n_feats,3,act_type='lrelu')



        #self.upsampling  = nn.PixelShuffle(2)


        # ###############
        # in_channels = int(n_feats//4)
        #
        # self.distilled_channels = int(in_channels * distillation_rate)
        # self.remaining_channels = int(in_channels - self.distilled_channels)
        # #self.c1 = conv_layer(in_channels, in_channels, 3)
        # self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        #self.act = activation('lrelu', neg_slope=0.01)
       # self.act2 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)
      #  self.act3 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)

        # self.scale1 = Scale(1)
        # self.scale2 = Scale(1)
        # self.scale3 = Scale(1)
        # self.scale4 = Scale(1)
        ###############

        self.bottom11 = conv_block(n_feats,n_feats,1,act_type=None)
        self.bottom11_dw = conv_block(n_feats, n_feats,5, groups=n_feats,act_type=None)

        self.reduction = conv_block((n_feats // reduction_factor)*3, n_feats // reduction_factor,1,act_type=None)
    def forward(self, x):


        x  = self.conv00(self.conv01(x))
        rc = self.reduce_channels(x)

        rs = self.reduce_spatial_size(rc)






        pool = self.pool(rs)

      #  print("After ", pool.shape)
        conv1 = self.conv1(pool)
        conv2 = self.conv2(conv1)+conv1
        conv3 = self.conv2(conv2)+conv2

        conv = self.reduction(torch.cat([conv1,conv2,conv3],1))




      #  print(rc.shape)

       # print(up.shape)

        up =  torch.nn.functional.upsample(conv, size=(rc.shape[2],rc.shape[3]), mode='nearest')
        up =  up + rc
        out = (self.sigmoid(self.increase_channels(up)) * x) *  self.sigmoid(self.bottom11_dw(self.bottom11(x)))





        #out_c1 = self.act(self.c1(up))

        # distilled_c1, remaining_c1 = torch.split(up, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c2 = self.act(self.c2(remaining_c1))
        # distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c3 = self.act(self.c3(remaining_c2))
        # distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c4 = self.c4(remaining_c3)

        # F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        # F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))
        #
        # out = self.fusion3(torch.cat([F1, F2],1))





  #       out = self.increase_channels(torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1))
  #       out = self.reduce_spatial_size (out) #+ up
  #      #
  #
  # #       rc = self.reduce_spatial_size(rc)
  #
  #       out = self.sigmoid(out)*x


        return out

class ESAB_MODIFIED(nn.Module):
    def __init__(self, n_feats, reduction_factor=4, distillation_rate=0.25):
        super(ESAB_MODIFIED, self).__init__()



        self.reduce_channels = nn.Conv2d(n_feats, n_feats//reduction_factor,1)
        self.reduce_spatial_size = nn.Conv2d(n_feats//reduction_factor, n_feats//reduction_factor, 3, stride=2, padding=1)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//reduction_factor, n_feats, 1)

        self.conv1 = conv_block(n_feats//reduction_factor, n_feats//reduction_factor,3,dilation=1,act_type='lrelu')
        self.conv2 = conv_block(n_feats // reduction_factor, n_feats//reduction_factor, 3,dilation=2,act_type='lrelu')
        #self.conv3 = conv_block(n_feats // 4, n_feats//4, 3)

        self.sigmoid = nn.Sigmoid()


        self.conv00 = conv_block(n_feats, n_feats,3, act_type=None)
        self.conv01 = conv_block(n_feats, n_feats,3,act_type='lrelu')



        #self.upsampling  = nn.PixelShuffle(2)


        # ###############
        # in_channels = int(n_feats//4)
        #
        # self.distilled_channels = int(in_channels * distillation_rate)
        # self.remaining_channels = int(in_channels - self.distilled_channels)
        # #self.c1 = conv_layer(in_channels, in_channels, 3)
        # self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        # self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        #self.act = activation('lrelu', neg_slope=0.01)
       # self.act2 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)
      #  self.act3 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)

        self.scale1 = Scale(1)
        self.scale2 = Scale(1)
        self.scale3 = Scale(1)
        self.scale4 = Scale(1)
        ###############

        self.bottom11 = conv_block(n_feats,n_feats,1,act_type=None)
        self.bottom11_dw = conv_block(n_feats, n_feats,5, groups=n_feats,act_type=None)



    def forward(self, x):


        x  = self.conv00(self.conv01(x))
        rc = self.reduce_channels(x)

        rs = self.reduce_spatial_size(rc)






        pool = self.pool(rs)

      #  print("After ", pool.shape)
        conv = self.conv2(pool)
        conv = conv + self.conv1(pool)


      #  print(rc.shape)

       # print(up.shape)

        up =  torch.nn.functional.upsample(conv, size=(rc.shape[2],rc.shape[3]), mode='nearest')
        up =  up + rc
        out = (self.sigmoid(self.increase_channels(up)) * x) *  self.sigmoid(self.bottom11_dw(self.bottom11(x)))





        #out_c1 = self.act(self.c1(up))

        # distilled_c1, remaining_c1 = torch.split(up, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c2 = self.act(self.c2(remaining_c1))
        # distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c3 = self.act(self.c3(remaining_c2))
        # distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        # out_c4 = self.c4(remaining_c3)

        # F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        # F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))
        #
        # out = self.fusion3(torch.cat([F1, F2],1))





  #       out = self.increase_channels(torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1))
  #       out = self.reduce_spatial_size (out) #+ up
  #      #
  #
  # #       rc = self.reduce_spatial_size(rc)
  #
  #       out = self.sigmoid(out)*x


        return out

class ESAB(nn.Module):
    def __init__(self, n_feats, reduction_factor=4, distillation_rate=0.25):
        super(ESAB, self).__init__()



        self.reduce_channels = nn.Conv2d(n_feats, n_feats//reduction_factor,1)
        self.reduce_spatial_size = nn.Conv2d(n_feats//reduction_factor, n_feats//reduction_factor, 3, stride=2, padding=1)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//reduction_factor, n_feats, 1)

        self.conv1 = conv_block(n_feats//reduction_factor, n_feats//reduction_factor,3,dilation=1,act_type='lrelu')
        self.conv2 = conv_block(n_feats // reduction_factor, n_feats//reduction_factor, 3,dilation=2,act_type='lrelu')

        self.sigmoid = nn.Sigmoid()


        self.conv00 = conv_block(n_feats, n_feats,3)
        self.conv01 = conv_block(n_feats, n_feats,3,act_type='lrelu')



        self.act = activation('lrelu', neg_slope=0.01)

        ###############




    def forward(self, x):

        x  = self.conv00(self.conv01(x))
        rc = self.reduce_channels(x)

        rs = self.reduce_spatial_size(rc)






        pool = self.pool(rs)

        conv = self.conv2(pool)
        conv = conv + self.conv1(pool)


        up =  torch.nn.functional.upsample(conv, size=(rc.shape[2],rc.shape[3]), mode='nearest')
        up =  up + rc
        out = self.sigmoid(self.increase_channels(up)) * x



        return out



class ESALayer_modified(nn.Module):
    def __init__(self, n_feats, distillation_rate=0.25):
        super(ESALayer_modified, self).__init__()



        self.reduce_channels = conv_block(n_feats, n_feats,3)
        self.reduce_spatial_size = conv_block(n_feats, n_feats, 3, stride=2)
       # self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//4, n_feats, 3)

        # self.conv1 = conv_block(n_feats//4, n_feats//4,3)
        # self.conv2 = conv_block(n_feats // 4, n_feats//4, 3)
        # self.conv3 = conv_block(n_feats // 4, n_feats//4, 3)

        self.sigmoid = nn.Sigmoid()


        self.upsampling  = nn.PixelShuffle(2)


        ###############
        in_channels = int(n_feats//4)

        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        #self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 5)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 5)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 5)
        self.act = nn.LeakyReLU(inplace=True)# activation('lrelu', neg_slope=0.05)
       # self.act2 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)
      #  self.act3 = nn.LeakyReLU(inplace=True)  # activation('lrelu', neg_slope=0.05)

        self.scale1 = Scale(1)
        self.scale2 = Scale(1)
        self.scale3 = Scale(1)
        self.scale4 = Scale(1)
        ###############


    def forward(self, x):

      #  rc = (self.reduce_channels(x))



        up = self.upsampling(x)

        #pool = self.pool(rs)

        #conv = self.conv1(self.conv2(self.conv3(pool)))




        #out_c1 = self.act(self.c1(up))

        distilled_c1, remaining_c1 = torch.split(up, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)

        # F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        # F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))
        #
        # out = self.fusion3(torch.cat([F1, F2],1))





        out = self.increase_channels(torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1))
        out = self.reduce_spatial_size (out) #+ up
       #

  #       rc = self.reduce_spatial_size(rc)

        out = self.sigmoid(out)*x


        return (out, out+x)



class ESAB_MODIFIED_xunit(nn.Module):
    def __init__(self, n_feats, reduction_factor=4, distillation_rate=0.25):
        super(ESAB_MODIFIED_xunit, self).__init__()
        self.reduce_channels = nn.Conv2d(n_feats, n_feats//reduction_factor,1)
        self.reduce_spatial_size = nn.Conv2d(n_feats//reduction_factor, n_feats//reduction_factor, 3, stride=2, padding=1)
        self.pool = nn.MaxPool2d(7, stride=3)
        self.increase_channels = conv_block(n_feats//reduction_factor, n_feats, 1)
        self.conv1 = Modulecell_LR(n_feats//reduction_factor, n_feats//reduction_factor,3,dilation=1)
        self.conv2 = Modulecell_LR(n_feats // reduction_factor, n_feats//reduction_factor, 3,dilation=2)
        self.sigmoid = nn.Sigmoid()
        self.conv00 = conv_block(n_feats, n_feats,3, act_type=None)
        self.conv01 = Modulecell_LR(n_feats, n_feats,3)
        self.bottom11 = conv_block(n_feats,n_feats,1,act_type=None)
        self.bottom11_dw = conv_block(n_feats, n_feats,5, groups=n_feats,act_type=None)
    """
        self.conv1 = Modulecell_LR(n_feats//reduction_factor, n_feats//reduction_factor,3,dilation=1)
        self.conv2 = Modulecell_LR(n_feats // reduction_factor, n_feats//reduction_factor, 3,dilation=2)
        self.sigmoid = nn.Sigmoid()
        self.conv00 = Modulecell_LR(n_feats, n_feats,3)
    """
    def forward(self, x):
        x  = self.conv00(self.conv01(x))
        rc = self.reduce_channels(x)
        rs = self.reduce_spatial_size(rc)
        pool = self.pool(rs)
        conv = self.conv2(pool)
        conv = conv + self.conv1(pool)
        up =  torch.nn.functional.upsample(conv, size=(rc.shape[2],rc.shape[3]), mode='nearest')
        up =  up + rc
        out = (self.sigmoid(self.increase_channels(up)) * x) *  self.sigmoid(self.bottom11_dw(self.bottom11(x)))
        return out

class Modulecell_LR(nn.Module):
    def __init__(self,in_channels=1,out_channels=64,kernel_size=3,skernel_size=9, dilation = 1):
        super(Modulecell_LR,self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(in_channels,out_channels,kernel_size=kernel_size,padding=int((kernel_size - 1) / 2) * dilation,bias=True, dilation = dilation))
        self.module = nn.Sequential(
            nn.LeakyReLU(0.05,1),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, groups=out_channels),
            nn.Conv2d(out_channels, out_channels, kernel_size=1),
            Gaussian())
    def forward(self,x):
        x1 = self.features(x)
        x2 = self.module(x1)
        x = torch.mul(x1,x2)
        return x

class Gaussian(nn.Module):
    def forward(self,input):
        return torch.exp(-torch.mul(input,input))



class Scale(nn.Module):
    def __init__(self, init_value=1e-3):
        super().__init__()
        self.scale = nn.Parameter(torch.FloatTensor([init_value]))

    def forward(self, input):
        return input * self.scale



class IMDModule_RFA(nn.Module):
    def __init__(self, in_channels, distillation_rate=0.25):
        super(IMDModule_RFA, self).__init__()
        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        #self.c5 = conv_layer(in_channels, in_channels, 1)
        self.cca = CCALayer(self.distilled_channels*4)



        self.fusion1 = conv_layer(self.distilled_channels*2, self.distilled_channels, 1)
        self.fusion2 = conv_layer(self.distilled_channels*2, self.distilled_channels, 1)
        self.fusion3 = conv_layer(self.distilled_channels*2, self.distilled_channels*4, 1)


    def forward(self, input):
        out_c1 = self.act(self.c1(input))
        distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)

        F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))

        out = self.fusion3(torch.cat([F1, F2],1))





        #out = torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1)

        out_fused = self.cca(out) + input
        return out_fused


class IMDModule_smaller_bff_ESA(nn.Module):
    def __init__(self, in_channels, distillation_rate=0.25):
        super(IMDModule_smaller_bff_ESA, self).__init__()

        self.shuffle = nn.PixelShuffle(2)

        in_channels = in_channels // 4

        self.deshuffle = conv_block(in_channels, in_channels * 4, 3, 2)



        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        #self.c5 = conv_layer(in_channels, in_channels, 1)
        #self.cca = CCALayer(self.distilled_channels*4)
        self.cca = CCALayer(self.distilled_channels * 4  *4 , reduction=16)


        #
        # self.fusion1 = conv_layer(self.distilled_channels*2, self.distilled_channels, 1)
        # self.fusion2 = conv_layer(self.distilled_channels*2, self.distilled_channels, 1)
        # self.fusion3 = conv_layer(self.distilled_channels*2, self.distilled_channels*4, 1)
        #
        #
        self.scale1 = Scale(1)
        self.scale2 = Scale(1)
        self.scale3 = Scale(1)
        self.scale4 = Scale(1)
    def forward(self, input):


        input_new  = self.shuffle(input)

        out_c1 = self.act(self.c1(input_new))
        distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)

        # F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        # F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))
        #
        # out = self.fusion3(torch.cat([F1, F2],1))





        out = torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1)

        out = self.deshuffle(out)
        out_fused = self.cca(out) + input
        return out_fused


class IMDModule_smaller_bff(nn.Module):
    def __init__(self, in_channels, distillation_rate=0.25):
        super(IMDModule_smaller_bff, self).__init__()
        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        #self.c5 = conv_layer(in_channels, in_channels, 1)
        self.cca = CCALayer(self.distilled_channels*4)



        self.fusion1 = conv_layer(self.distilled_channels*2, self.distilled_channels, 1)
        self.fusion2 = conv_layer(self.distilled_channels*2, self.distilled_channels, 1)
        self.fusion3 = conv_layer(self.distilled_channels*2, self.distilled_channels*4, 1)


    def forward(self, input):
        out_c1 = self.act(self.c1(input))
        distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)

        F1 = self.fusion1 (torch.cat([distilled_c1, distilled_c2],1))
        F2 = self.fusion2(torch.cat([distilled_c3, out_c4],1))

        out = self.fusion3(torch.cat([F1, F2],1))





        #out = torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1)

        out_fused = self.cca(out) + input
        return out_fused

class IMDModule_smallest(nn.Module):
    def __init__(self, in_channels, reduction=16, distillation_rate=0.5):
        super(IMDModule_smallest, self).__init__()
        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
   #     self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
   #     self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        self.c5 = conv_layer(in_channels, in_channels, 1)
        self.cca = CCALayer(self.distilled_channels * 2, reduction=reduction)

       # self.scale1 = Scale(1)
      #  self.scale2 = Scale(1)
        self.scale3 = Scale(1)
        self.scale4 = Scale(1)

    def forward(self, input):
        out_c1 = self.act(self.c1(input))
       # distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
       # out_c2 = self.act(self.c2(remaining_c1))
        #distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
      #  out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)
        out = torch.cat([self.scale3(distilled_c3), self.scale4(out_c4)], dim=1)
        out_fused = self.c5(self.cca(out)) + input
        return out_fused


class IMDModule_smaller(nn.Module):
    def __init__(self, in_channels, reduction=16, distillation_rate=0.25):
        super(IMDModule_smaller, self).__init__()
        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        self.c5 = conv_layer(in_channels, in_channels, 1)
        self.cca = CCALayer(self.distilled_channels * 4, reduction=reduction)

        self.scale1 = Scale(1)
        self.scale2 = Scale(1)
        self.scale3 = Scale(1)
        self.scale4 = Scale(1)

    def forward(self, input):
        out_c1 = self.act(self.c1(input))
        distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)
        out = torch.cat([self.scale1(distilled_c1), self.scale2(distilled_c2), self.scale3(distilled_c3), self.scale4(out_c4)], dim=1)
        out_fused = self.c5(self.cca(out)) + input
        return out_fused


class IMDModule(nn.Module):
    def __init__(self, in_channels, distillation_rate=0.25):
        super(IMDModule, self).__init__()
        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        self.c5 = conv_layer(in_channels, in_channels, 1)
        self.cca = CCALayer(self.distilled_channels * 4)

    def forward(self, input):
        out_c1 = self.act(self.c1(input))
        distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)
        out = torch.cat([distilled_c1, distilled_c2, distilled_c3, out_c4], dim=1)
        out_fused = self.c5(self.cca(out)) + input
        return out_fused

class IMDModule_speed(nn.Module):
    def __init__(self, in_channels, distillation_rate=0.25):
        super(IMDModule_speed, self).__init__()
        self.distilled_channels = int(in_channels * distillation_rate)
        self.remaining_channels = int(in_channels - self.distilled_channels)
        self.c1 = conv_layer(in_channels, in_channels, 3)
        self.c2 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c3 = conv_layer(self.remaining_channels, in_channels, 3)
        self.c4 = conv_layer(self.remaining_channels, self.distilled_channels, 3)
        self.act = activation('lrelu', neg_slope=0.05)
        self.c5 = conv_layer(self.distilled_channels * 4, in_channels, 1)

    def forward(self, input):
        out_c1 = self.act(self.c1(input))
        distilled_c1, remaining_c1 = torch.split(out_c1, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c2 = self.act(self.c2(remaining_c1))
        distilled_c2, remaining_c2 = torch.split(out_c2, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c3 = self.act(self.c3(remaining_c2))
        distilled_c3, remaining_c3 = torch.split(out_c3, (self.distilled_channels, self.remaining_channels), dim=1)
        out_c4 = self.c4(remaining_c3)

        out = torch.cat([distilled_c1, distilled_c2, distilled_c3, out_c4], dim=1)
        out_fused = self.c5(out) + input
        return out_fused

def pixelshuffle_block(in_channels, out_channels, upscale_factor=2, kernel_size=3, stride=1):
    conv = conv_layer(in_channels, out_channels * (upscale_factor ** 2), kernel_size, stride)
    pixel_shuffle = nn.PixelShuffle(upscale_factor)
    return sequential(conv,pixel_shuffle)
